#!/usr/bin/env python3
import rospy
from sensor_msgs.msg import Image, CompressedImage, CameraInfo
from cv_bridge import CvBridge
import cv2
import os
import time
import yaml

def load_camera_info():
    # Path to the YAML file
    yaml_path = "~/utils_ws/src/camera_sdk/param/param.yaml"
    
    try:
        with open(yaml_path, 'r') as file:
            params = yaml.safe_load(file)
            
        # Create CameraInfo message
        camera_info = CameraInfo()
        
        # Set header (frame_id typically matches camera name)
        camera_info.header.frame_id = params['camera_name']
        
        # Set image dimensions
        camera_info.width = params['image_width']
        camera_info.height = params['image_height']
        
        # Set camera matrix (flatten 3x3 matrix)
        camera_info.K = params['camera_matrix']['data']
        
        # Set distortion model and coefficients
        camera_info.distortion_model = params['distortion_model']
        camera_info.D = params['distortion_coefficients']['data']
        
        # Set rectification matrix (flatten 3x3 matrix)
        camera_info.R = params['rectification_matrix']['data']
        
        # Set projection matrix (flatten 3x4 matrix)
        camera_info.P = params['projection_matrix']['data']
        
        return camera_info
        
    except Exception as e:
        rospy.logerr(f"Failed to load camera parameters: {e}")
        return None

def publish_camera_feed():
    # 初始化 ROS 节点
    rospy.init_node('mjpeg_camera_publisher', anonymous=True)
    
    # 创建图像发布者
    image_pub = rospy.Publisher('/camera/image_raw', Image, queue_size=10)
    # 创建压缩图像发布者
    compressed_pub = rospy.Publisher('/camera/image_raw/compressed', CompressedImage, queue_size=10)
    # 创建相机信息发布者
    camera_info_pub = rospy.Publisher('/camera/camera_info', CameraInfo, queue_size=10)
    
    # 加载相机内参
    camera_info = load_camera_info()
    if camera_info is None:
        rospy.logerr("无法加载相机参数，继续运行但不会发布相机信息")
    
    # 初始化 OpenCV 到 ROS 的转换桥
    bridge = CvBridge()

    # 图像保存参数
    save_images = False  # 是否保存图像的布尔标志
    save_frequency = 3  # 默认保存频率 (Hz)
    save_interval = 1.0 / save_frequency  # 保存间隔 (秒)
    last_save_time = time.time()
    save_dir = "~/utils_ws/src/camera_sdk/image/"  # 保存目录
    
    # 创建保存目录
    if save_images and not os.path.exists(save_dir):
        os.makedirs(save_dir)

    # 打开摄像头
    cap = cv2.VideoCapture('/dev/video0', cv2.CAP_V4L2)
    cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'MJPG'))
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)  
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)  

    # 检查摄像头是否成功打开
    if not cap.isOpened():
        rospy.logerr("无法打开摄像头")
        return

    rate = rospy.Rate(100)  # 设置发布频率
    while not rospy.is_shutdown():
        ret, frame = cap.read()
        if not ret:
            rospy.logwarn("无法读取摄像头图像")
            continue
        
        # 检查图像形状
        if frame is None:
            rospy.logwarn("获取的帧为空")
            continue
        rospy.loginfo(f"Frame shape: {frame.shape}")  # 打印图像形状

        try:
            # 获取当前时间
            current_time = rospy.Time.now()
            
            # 发布原始图像
            img_msg = bridge.cv2_to_imgmsg(frame, encoding="bgr8")
            img_msg.header.stamp = current_time
            image_pub.publish(img_msg)
            
            # 发布压缩图像
            compressed_msg = CompressedImage()
            compressed_msg.header.stamp = current_time
            compressed_msg.format = "jpeg"
            # 使用cv2.imencode将图像压缩为JPEG格式
            ret, compressed_data = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
            if ret:
                compressed_msg.data = compressed_data.tobytes()
                compressed_pub.publish(compressed_msg)
            else:
                rospy.logwarn("图像压缩失败")
            
            # 发布相机信息
            if camera_info is not None:
                camera_info.header.stamp = current_time
                camera_info_pub.publish(camera_info)
            
            # 图像保存功能
            if save_images:
                current_save_time = time.time()
                if current_save_time - last_save_time >= save_interval:
                    # 生成带时间戳的文件名
                    timestamp = current_time.to_sec()
                    filename = os.path.join(save_dir, f"image_{timestamp:.3f}.jpg")
                    cv2.imwrite(filename, frame)
                    rospy.loginfo(f"已保存图像: {filename}")
                    last_save_time = current_save_time
                
        except Exception as e:
            rospy.logerr(f"图像处理出错: {e}")

        rate.sleep()
    
    cap.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    try:
        publish_camera_feed()
    except rospy.ROSInterruptException:
        pass